#include <windows.h>
#include <stdio.h>
#include <stdlib.h>
#include "..\..\GetPEInfo\GetPEInfo\FileDialogs.h"
#include "WriteTestPEFile.h"
#include "WritePEFile.h"

int main(int argc, char** argv)
{
	char buff[0x100];
	TCHAR* filePath = NULL;
	BOOL freeFilePath = FALSE;
	if (argc == 1)
	{
		static const TCHAR  catagory [] = TEXT("Executable File");
		static const TCHAR* exts [] = {TEXT("exe"), NULL};

SelectFileWithDialog:
		if (!(filePath = SaveFileDlg(catagory, exts, NULL)))
			return -1;
	}
	else
	{
		int i;
		for (i = 1; i <= argc; ++i)
		{
			if (GetFileAttributesA(argv[i]) != INVALID_FILE_ATTRIBUTES)
			{
#ifndef UNICODE
				filePath = argv[i];
#else
				unsigned int argvLen, j;

				freeFilePath = TRUE;
				argvLen = strlen(argv[i]);
				filePath = (TCHAR*)malloc(sizeof(TCHAR) * argvLen);
				for (j = 0; j < argvLen; ++j)
					filePath[j] = (TCHAR)argv[i][j];
				filePath[j] = (TCHAR)'\0';
#endif
			}
		}

		if (filePath == NULL)
			goto SelectFileWithDialog;
	}

#if 1
	{
		static const BYTE code[PE_PAGE_SIZE] =
		{
			0xB8, 0xFF, 0xFF, 0xFF, 0xFF, // MOV EAX, 0xFFFFFFFF
			0x40, // INC EAX
			0x83, 0xf8, 0xff, // CMP EAX, -1
			0x75, 0xfa, 0x90, // JNE -2
			0xc3 // RET
		};

		FILE* file;
		struct PEFileCreationParams p;
		struct PEFileSection textSect;

		textSect.attr = PFSAF_EXECUTE | PFSAF_READ;
		textSect.numPEPages = 1;
		textSect.data = code;

		p.machineType = MT_IA32;
		p.numSections = 1;
		p.sections = (struct PEFileSection**)malloc(sizeof(struct PEFileSection*));
		*p.sections = &textSect;

		// Open file
#ifndef UNICODE
		fopen_s(&file, filePath, "wb");
#else
		_wfopen_s(&file, filePath, L"wb");
#endif

		WritePEFile(file, &p);

		fclose(file);
	}
#else
	WriteTestPEFile(filePath);
#endif
	sprintf(buff, "C:\\Software\\OllyDbg\\OLLYDBG.exe  \"%s\"", filePath);
	system(buff);
	system("pause");

	if (freeFilePath)
		free(filePath);

	return 0;
}